import os
from pathlib import Path

import cv2
import numpy as np
import pytesseract
from Sudoku.DigitRecognizer import DigitRecognizer


class Png2:
    def __init__(self, tesseract_path):
        """
        初始化 SudokuSolver 类。
        :param tesseract_path: Tesseract OCR 的可执行文件路径
        """
        pytesseract.pytesseract.tesseract_cmd = tesseract_path
        self.digit_recognizer = DigitRecognizer()

    def preprocess_image(self, img):
        """图像预处理：灰度化和二值化"""
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)  # 灰度化
        # 使用自适应阈值进行二值化处理，改善图像内容的提取
        thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                       cv2.THRESH_BINARY_INV, 11, 2)
        return thresh

    def extract_largest_contour(self, img):
        """提取最大轮廓区域"""
        contours, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        # 找到面积最大的轮廓
        largest = max(contours, key=cv2.contourArea)
        mask = np.zeros_like(img)
        # 在掩码上绘制最大轮廓（填充轮廓内部）
        cv2.drawContours(mask, [largest], -1, (255, 255, 255), thickness=cv2.FILLED)
        # 使用掩码提取轮廓内的区域
        extracted = cv2.bitwise_and(img, img, mask=mask)
        return extracted, largest

    def cut_image_into_grid(self, img, points, grid_size=(9, 9)):
        """将绿框内的图像按指定网格大小切割"""
        if points is None:
            return []

        result = img.copy()
        cv2.drawContours(result, [points], -1, (0, 255, 0), 3)

        # 获取四边形的外接矩形
        rect = cv2.boundingRect(points)
        x, y, w, h = rect

        # 裁剪出绿框内的区域
        cropped_img = img[y:y + h, x:x + w]

        # 计算每个网格的宽度和高度
        grid_width = w // grid_size[0]
        grid_height = h // grid_size[1]

        # 存储切割后的小块
        grid_images = []

        # 按照网格大小切割
        for i in range(grid_size[0]):
            for j in range(grid_size[1]):
                start_x = i * grid_width
                start_y = j * grid_height
                end_x = (i + 1) * grid_width
                end_y = (j + 1) * grid_height

                # 裁剪每个网格区域
                grid_images.append(cropped_img[start_y:end_y, start_x:end_x])

        return grid_images

    def recognize_digits_from_images(self, image_list):
        """识别小图中的数字"""
        recognized_digits = []
        for idx, image in enumerate(image_list):
            # 将小图转换为PIL格式
            height, width = image.shape[:2]
            # 计算切割后的区域，去掉四周x个像素
            cropped_image = image[8:height - 8, 9:width - 8]
            # 调用识别接口
            digit = self.digit_recognizer.recognize_digit(cropped_image)
            recognized_digits.append(digit[0])
        return recognized_digits

    def solve_sudoku(self, image_path):
        """
        :param image_path: 数独图像路径
        :return: 识别出的数字列表
        """
        # 读取图像
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError("无法加载图像，请检查路径是否正确。")

        # 图像预处理
        processed = self.preprocess_image(img)

        # 提取最大轮廓区域
        extracted, largest_contour = self.extract_largest_contour(processed)

        # 切割图像
        grid_images = self.cut_image_into_grid(extracted, largest_contour)

        self.save_grid_images(grid_images)

        # 识别切割后图片中的数字
        recognized_digits = self.recognize_digits_from_images(grid_images)

        return recognized_digits

    def save_grid_images(self, grid_images, output_folder="output_grids"):
        """
        将切割后的图片保存到本地文件夹中

        :param grid_images: 切割后的图片列表
        :param output_folder: 保存图片的文件夹路径
        """
        # 创建输出文件夹（如果不存在）
        Path(output_folder).mkdir(parents=True, exist_ok=True)

        # 遍历图片列表并保存
        for idx, image in enumerate(grid_images):
            if image is not None:  # 确保图片不为空
                file_path = os.path.join(output_folder, f"grid_{idx + 1}.png")
                cv2.imwrite(file_path, image)
            else:
                print(f"第 {idx + 1} 张图片为空，跳过保存")


# 示例用法
if __name__ == '__main__':
    """
    调用DigitRecognizer接口，图片转换和识别出数字
    传入截图，自动分割和识别
    """
    # 初始化 SudokuSolver
    solver = Png2(tesseract_path=r'tesseract.exe')
    # 传入截图，自动分割和识别
    image_path = "png.jpg"  # 替换为你的图像路径
    try:
        digits = solver.solve_sudoku(image_path)
        print("识别出的数字：", digits)
    except Exception as e:
        print("发生错误：", str(e))
